# Skecth generation method 2
import random
from itertools import chain

import torch
from PIL import Image
from scipy.ndimage import gaussian_filter
from skimage.filters import threshold_otsu
import numpy as np
import os
import cv2
from tqdm import tqdm
import torchvision.transforms as T
import torchvision.transforms.functional as TF
from torch.utils import data
from main import get_config
from utils.FileHelper import openRGBImage, isImage, listdir, toNormalPic
import matplotlib.pyplot as plt


class Jitter:
    def __init__(self, brightness_factor, contrast_factor, saturation_factor, hue_factor):
        self.brightness_factor = brightness_factor
        self.contrast_factor = contrast_factor
        self.saturation_factor = saturation_factor
        self.hue_factor = hue_factor

    def __call__(self, img, *args, **kwargs):
        x = TF.adjust_brightness(img, self.brightness_factor)
        x = TF.adjust_contrast(x, self.contrast_factor)
        x = TF.adjust_saturation(x, self.saturation_factor)
        x = TF.adjust_hue(x, self.hue_factor)
        return x


def generate_lineart(img_path, desc_dir, lineart_generator):
    *_, artist, img = img_path.split("/")
    folder = os.path.join(desc_dir, artist).replace('\\', '/')
    os.makedirs(folder, exist_ok=True)

    sketch_name = f"{img.split('.')[-2]}.png"
    path = os.path.join(folder, sketch_name).replace('\\', '/')

    if os.path.isfile(path):
        return path

    img = openRGBImage(img_path)
    img_ndarray = np.array(img) / 255.0
    sketch = (cv2.cvtColor(lineart_generator(img_ndarray).astype(np.float32), cv2.COLOR_GRAY2RGB) * 255).astype(
        np.int32)
    cv2.imwrite(path, sketch)
    return path


class xDog:
    def __init__(self,
                 γ: float = 0.95,
                 ϕ: float = 1e9,
                 ϵ: float = -1e1,
                 k: float = 4.5,
                 σ: float = 0.3, ) -> None:
        self.γ = γ
        self.ϕ = ϕ
        self.ϵ = ϵ
        self.k = k
        self.σ = σ

    def __call__(self, img) -> np.ndarray:
        x = (img[..., 0] + img[..., 1] + img[..., 2]) / 3

        gaussian_a = gaussian_filter(x, self.σ)
        gaussian_b = gaussian_filter(x, self.σ * self.k)

        dog = gaussian_a - self.γ * gaussian_b

        inf = dog < self.ε
        xdog = inf * 1 + ~inf * (1 - np.tanh(self.φ * dog))

        xdog -= xdog.min()
        xdog /= xdog.max()
        xdog = xdog >= threshold_otsu(xdog)
        xdog = 1 - xdog

        return xdog


def make_dataset(ImageFolder: str, sketchFolder=None):
    if not os.path.isdir(ImageFolder):
        raise Exception("please provide folder of Image")
    images_list = chain(*(filter(isImage, listdir(f)) for f in listdir(ImageFolder)))
    images_list = sorted(list(images_list))

    # if sketchFolder is not None:
    #     sketch_list = chain(*(filter(isimg, listdir(f)) for f in listdir(sketchFolder)))
    #     sketch_list = sorted(list(sketch_list))
    #     return images_list, sketch_list

    if sketchFolder is None:
        sketch_dirname = os.path.abspath(os.path.join(ImageFolder, "..", 'sketches')).replace('\\', '/')
    else:
        sketch_dirname = sketchFolder
    if not os.path.isdir(sketch_dirname):
        os.makedirs(sketch_dirname)
    # sketchs_root = os.path.join(sketch_dirname, folderName).replace('\\', '/')
    pbar = tqdm(images_list)
    line_generator = xDog()
    sketch_list = list()
    for img in pbar:
        path = generate_lineart(img, sketch_dirname, line_generator)
        sketch_list.append(path)
    sketch_list = sorted(sketch_list)
    return images_list, sketch_list


class RandomHorizontalFlip(T.RandomHorizontalFlip):
    def __call__(self, img):
        hflip = np.random.rand() > 0.5
        imgs = TF.hflip(img) if hflip else img
        return imgs


class RandomVerticalFlip(T.RandomVerticalFlip):
    def __call__(self, img):
        vflip = np.random.rand() > 0.5
        imgs = TF.vflip(img) if vflip else img
        return imgs


class Dataset(data.Dataset):
    def __init__(self, config):
        root = config['TRAINING_CONFIG']['ROOT']
        self.contour_dir = os.path.join(root, 'contour')
        self.img_dir = os.path.join(root, 'img')
        labels = torch.load(os.path.join(root, 'labels.pt'))
        self.labels = labels['labels']
        self.groups = labels['groups']
        self.idxs = list(range(len(self.labels)))
        self.size = len(self.idxs)
        self.image_paths = [os.path.join(self.img_dir, '{:06d}.png'.format(i)) for i in self.idxs]
        self.contour_paths = [os.path.join(self.contour_dir, '{:06d}.png'.format(i)) for i in self.idxs]
        self.aug = config['TRAINING_CONFIG']['AUGMENT']
        self.noise_type = config['TRAINING_CONFIG']['NOISE']
        self.img_size = config['MODEL_CONFIG']['IMAGE_SIZE']

        if 'color_jitter' in self.aug:
            brightness_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            contrast_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            saturation_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            hue_factor = torch.tensor(1.0).uniform_(-.2, .2).item()
            self.jitter = Jitter(brightness_factor, contrast_factor, saturation_factor, hue_factor)
        if 'hflip' in self.aug:
            self.hflip = np.random.rand() > 0.5
        if 'vflip' in self.aug:
            self.vflip = np.random.rand() > 0.5
        if 'rotate' in self.aug:
            self.angle = T.RandomRotation.get_params([0, 360])

    def transform_style_image(self, x):
        if 'color_jitter' in self.aug:
            x = self.jitter(x)
        if 'hflip' in self.aug:
            x = TF.hflip(x) if self.hflip else x
        if 'vflip' in self.aug:
            x = TF.vflip(x) if self.vflip else x
        if 'rotate' in self.aug:
            x = TF.rotate(x, angle=self.angle)

        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
        return t(x)

    def transform_sketch(self, sketch):
        if 'hflip' in self.aug:
            sketch = TF.hflip(sketch) if self.hflip else sketch
        if 'vflip' in self.aug:
            sketch = TF.vflip(sketch) if self.vflip else sketch
        if 'rotate' in self.aug:
            sketch = TF.rotate(sketch, angle=self.angle)
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize(mean=0.5, std=0.5)])
        return t(sketch)

    def __len__(self):
        return self.size

    def __getitem__(self, idx1):
        label = self.labels[idx1]
        group = self.groups[label]

        # pick the icon in the same color cluster
        idx2 = random.choice(group)

        # pick another color cluster image
        idx3 = random.choice(self.idxs)

        s1 = Image.open(self.image_paths[idx1]).convert('RGB')
        s2 = Image.open(self.image_paths[idx2]).convert('RGB')
        s3 = Image.open(self.image_paths[idx3]).convert('RGB')
        contour = Image.open(self.contour_paths[idx3]).convert('L')

        s1 = self.transform_style_image(s1)
        s2 = self.transform_style_image(s2)
        s3 = self.transform_style_image(s3)
        contour = self.transform_sketch(contour)

        return s1, s2, s3, contour


# class Dataset(data.Dataset):
#
#     def __init__(self, config):
#         super(Dataset, self).__init__()
#         self.traindir = config['TRAINING_CONFIG']['TRAINING_DIR']
#         self.sketchdir = config['TRAINING_CONFIG']['SKETCH_DIR']
#         self.image_list, self.sketch_list = make_dataset(self.traindir, self.sketchdir)
#         if len(self.image_list) == 0 or len(self.sketch_list) == 0:
#             raise FileNotFoundError("found 0 images in folder")
#         self.xDog = xDog()
#         random.shuffle(self.image_list)
#         random.shuffle(self.sketch_list)
#
#         self.aug = config['TRAINING_CONFIG']['AUGMENT']
#         self.noise_type = config['TRAINING_CONFIG']['NOISE']
#
#         self.img_size = config['MODEL_CONFIG']['IMAGE_SIZE']
#         self.num_colors = config['TRAINING_CONFIG']['COLOR_NUMS']
#         self.pallete_size = int(math.sqrt(self.num_colors))
#
#         if 'gaussian' == self.noise_type:
#             self.mean = config['TRAINING_CONFIG']['MEAN']
#             self.std = config['TRAINING_CONFIG']['STD']
#         if 'color_jitter' in self.aug:
#             brightness_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
#             contrast_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
#             saturation_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
#             hue_factor = torch.tensor(1.0).uniform_(-.2, .2).item()
#             self.jitter = Jitter(brightness_factor, contrast_factor, saturation_factor, hue_factor)
#         if 'hflip' in self.aug:
#             self.hflip = RandomHorizontalFlip()
#         if 'vflip' in self.aug:
#             self.vflip = RandomVerticalFlip()
#
#     def transform_sketch(self, sketch):
#         if 'hflip' in self.aug:
#             sketch = self.hflip(sketch)
#         if 'vflip' in self.aug:
#             sketch = self.vflip(sketch)
#         t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize(mean=(0.5), std=(0.5))])
#         return t(sketch)
#
#     def transform_ref(self, ref):
#         if 'color_jitter' in self.aug:
#             ref = self.jitter(ref)
#         if 'hflip' in self.aug:
#             ref = self.hflip(ref)
#         if 'vflip' in self.aug:
#             ref = self.vflip(ref)
#
#         real_pallete = get_dominant_colors(ref, self.num_colors)
#         real_pallete = torch.Tensor(real_pallete)
#         real_pallete = real_pallete / 127.5 - 1
#
#         if self.noise_type == 'gaussian':
#             img_ = np.array(ref).copy()
#             img_ = img_ / 255.0
#             noise = np.random.normal(self.mean, self.std, img_.shape)
#             gaussian_out = img_ + noise
#             gaussian_out = np.clip(gaussian_out, 0, 1)
#             gaussian_out = np.uint8(gaussian_out * 255)
#             ref = Image.fromarray(gaussian_out).convert('RGB')
#
#         t1 = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
#         # t2 = T.Compose([T.Normalize((0.5,) * 3, (0.5,) * 3)])
#         # t2 = t2(real_pallete)
#         return t1(ref), real_pallete
#
#     def __getitem__(self, idx):
#         ref = self.image_list[idx]
#         sketch = self.sketch_list[idx]
#
#         ref = openRGBImage(ref)
#         sketch = Image.open(sketch).convert('L')
#         ref_sketch = (self.xDog(np.array(ref) / 255.0) * 255).astype(np.int32)
#         ref_sketch = Image.fromarray(ref_sketch).convert('L')
#
#         ref, real_pallete = self.transform_ref(ref)
#         sketch = self.transform_sketch(sketch)
#         ref_sketch = self.transform_sketch(ref_sketch)
#
#         return ref, sketch, real_pallete, ref_sketch
#
#     def __len__(self):
#         return len(self.image_list)


class EvaluationDataset(data.Dataset):
    def __init__(self, config):
        self.img_size = config['MODEL_CONFIG']['IMAGE_SIZE']
        self.sourceFolder = config['EVALUATION_CONFIG']['SOURECE_DIR']
        self.sketchFolder = config['EVALUATION_CONFIG']['SKETCH_DIR']
        self.referenceFolder = config['EVALUATION_CONFIG']['REFERENCE_DIR']
        self.image_list, self.sketch_list = make_dataset(self.sourceFolder, self.sketchFolder)
        refer_list = chain(*(filter(isImage, listdir(f)) for f in listdir(self.referenceFolder)))
        self.refer_list = sorted(list(refer_list))
        self.list = list(zip(self.sketch_list, self.refer_list))

    def transform_ref(self, img):
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
        return t(img)

    def transform_sketch(self, sketch):
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize(mean=(0.5), std=(0.5))])
        return t(sketch)

    def __getitem__(self, idx):
        sketch, ref = self.list[idx]
        sketch = Image.open(sketch).convert('L')
        ref = openRGBImage(ref)
        return self.transform_ref(ref), self.transform_sketch(sketch)

    def __len__(self):
        return len(self.refer_list)


def getTrainLoader(config):
    dataset = Dataset(config)
    dataloader = data.DataLoader(dataset=dataset,
                                 batch_size=config['TRAINING_CONFIG']['BATCH_SIZE'],
                                 shuffle=True,
                                 num_workers=config['TRAINING_CONFIG']['NUM_WORKER'],
                                 drop_last=True,
                                 pin_memory=True)
    return dataloader


def getEvaluationLoader(config):
    dataset = EvaluationDataset(config)
    evaluationLoader = data.DataLoader(dataset=dataset,
                                       batch_size=config['EVALUATION_CONFIG']['BATCH_SIZE'],
                                       shuffle=False,
                                       drop_last=False,
                                       pin_memory=True,
                                       num_workers=config['EVALUATION_CONFIG']['NUM_WORKER'])
    return evaluationLoader


def draw(colors):
    leng = len(colors)
    x = [xi for xi in range(0, leng + 1, 1)]
    y = [yi for yi in range(0, leng + 1, 1)]
    c = Image.new('RGB', (leng, leng))
    for i in range(leng):
        value = colors[i]
        c.putpixel([x[i], y[i]], tuple(value))
    c.save('./c.png')


if __name__ == '__main__':
    config = get_config()
    dataset = Dataset(config)
    s1, s2, s3, contour = dataset.__getitem__(0)
    s1 = toNormalPic(s1)
    s2 = toNormalPic(s2)
    s3 = toNormalPic(s3)
    contour = toNormalPic(contour)
    plt.subplot(2, 2, 1)
    plt.imshow(s1)
    plt.subplot(2, 2, 2)
    plt.imshow(s2)
    plt.subplot(2, 2, 3)
    plt.imshow(s3)
    plt.subplot(2, 2, 4)
    plt.imshow(contour, cmap='gray')
    plt.show()
